package mini.compiler;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.StringTokenizer;
import java.util.logging.Logger;

import mini.Proc;
import mini.Program;
import mini.util.MiniUtil;
import mini.util.SymbolicNames;


/**
 * Does stuff. See rewrite
 * 
 * @author Meyer
 */
public class ProgramRewriter {

	private Logger log = Logger.getLogger( "mini.compiler.ProgramRewriter" );
	
    /**
     * At this point, the input reader code refers to a lot of
     * variables that don't exist in global memory, because they
     * are local user and temp variables to some procedure. So
     * lines like:
     * 
     * LDA temp1
     * 
     * Have to be rewritten, according to slide 33, into four lines
     * that look like:
     * 
     * LDA FP
     * ADD <symbol>
     * STA FPB
     * LDI FBC
     * 
     * What's in the symbol is the offset to the variable in the
     * activation record. If you know what procedure you're in,
     * you can use the findSymbolicOffsetForVariable() method
     * on the current Proc object to see what <symbol> should be.
     * 
     * One trick is that you have to figure out what procedure
     * you're in. Procedures start with a line that looks like:
     * 
     * ProcedureN: LDA temp1      => e.g. Procedure2: LDA Constant0
     * 
     * You can use the string "ProcedureN" and compare it to the
     * names of the procedure objects (found with the getLocationLabel()
     * method on procedures) until you find one that matches. That name
     * CANNOT be used as a key into the functiontable map to find
     * the procedure as it's not the procedure's name, it's its label
     * 
     * @param hm Symbol table
     * @param functiontable Map of function names to their procedures
     * @param An open input reader to some symbolic, optimized code
     * @param An open output reader to write the rewritten code
     */
	public void rewrite(HashMap<String,Integer> hm, HashMap<String, Proc> functiontable,
			BufferedReader inputReader, BufferedWriter outputWriter) {

		try {
			// We should always be within some procedure's code at any point in time
			// That procedure is stored here
			Proc proc = null;
			
			int totalLineCounter = 0;	// This one includes counting comments and blank lines, so we can use it for error reporting
			
			// Step through the lines, one by one
			String line = inputReader.readLine();
			while( line != null ) {
				
				totalLineCounter++;
				
				if( line.trim().length() == 0 ) {
					// If it's a blank line, don't both to copy it
				}
				else if( line.startsWith(";") ) {
					// If it's a comment line, just copy it over
					outputWriter.write(line);
					outputWriter.write(Program.NEWLINE);
				} else {					
					proc = rewriteASingleLine(hm, functiontable, outputWriter,
							proc, totalLineCounter, line);
				}
				
				line = inputReader.readLine();
			}
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

	}

	/**
	 * This rewrites just a single, non-blank line in a larger file
	 * 
	 * @param hm Symbol table
	 * @param functiontable Funtion table
	 * @param outputWriter Where to write it
	 * @param proc Procedure we might be in
	 * @param totalLineCounter The line that we're on, number wise
	 * @param line The actual line
	 * @return A new procedure, if we happen to move to a new
	 * procedure within this line
	 * 
	 * @throws IOException
	 */
	private Proc rewriteASingleLine(HashMap<String,Integer> hm, HashMap<String, Proc> functiontable,
			BufferedWriter outputWriter, Proc proc, int totalLineCounter,
			String line) throws IOException {
		
		// First, remove any comments
		String commentlessLine = MiniUtil.removeAnyComments(line);
		int commentStartAt = line.indexOf(";");
		String commentStr = null;	// This line will hold comments to append at the end of the line
		if( commentStartAt != -1 ) {
			commentStr = line.substring(commentStartAt+1);
		}

		// Now, strip off any markers 
		String strippedLine = MiniUtil.removeAnyLocationMarkers(commentlessLine);
		String markerStr = null;	// This line will hold markers to prepend at the beginning of a line
		int markerIndex = line.indexOf(strippedLine);
		if( markerIndex > 0 ) {
			markerStr = line.substring(0,markerIndex);
		}

		// Are we in a new procedure?
		String[] markers = MiniUtil.getMarkersFromLine( line );
		if( markers != null ) {
			for( String marker : markers ) {
				if( SymbolicNames.isProcedureLabel(marker) ) {
					proc = findProcedureForMarker(functiontable,marker);
				}
			}
		}
		
		// Parse out the operator
		StringTokenizer strtok = new StringTokenizer(strippedLine);
		String instruction = null;
		if( strtok.hasMoreTokens() ) {
			instruction = strtok.nextToken();
		}
		
		// And maybe an operand with its symbolic offset
		String operand = null;
		String symbolicOffset = null;
		if( strtok.hasMoreTokens() ) {
			operand = strtok.nextToken();

			// See if we can find the operand within the procedure
			if( proc != null ) {
				symbolicOffset = proc.findSymbolicOffsetForVariable(operand);
				//log.info( "In procedure '" + proc.getLocationLabel() + "' the variable '" +
				//		operand + "' had a symbolic offset of '" + symbolicOffset + "' from the FP." );
			}
		}
		
		// We only care about rewriting some instructions
		if( potentialRewriteNeeded(instruction) && operand == null ) {
			
			// We need an operand if we have to rewrite it, though
			MiniUtil.reportParseError(totalLineCounter, line, "rewriting");
			
		} else if( potentialRewriteNeeded(instruction) && symbolicOffset == null) {
			
			// If we didn't find a symbolic offset, then hopefully it's a constant...
			if( !SymbolicNames.isConstant(operand) ) {
				System.err.println( "Probable error in Rewriter: Referencing '" + operand + "' as a global constant, though it does not appear to be based on naming conventions." );
			}
			// In which case we can just copy over the line
			outputWriter.write(line);
			outputWriter.write(Program.NEWLINE);
			
		} else {

			// We may have to translate this single line into multiple other lines
			String[] rewriteLines = null;
			
			// Figure out what we have to rewrite it as (potentially)
			if( CodeGenerator.LOAD_ACCUMULATOR_RELATIVE_TO_PREV_FP.equals(instruction) ) {
				// This is a meta instruction. It's like LDA, but instead of loading relative to
				// the address in FP, you load relative to the address in the scratch variable
				rewriteLines = rewriteLoadAccRelativeToPrevFPInstruction(commentStr, markerStr,	symbolicOffset);
			} else if( CodeGenerator.ADD_SIZE_OF_ACTIVATION_RECORD.equals(instruction) ) {
				rewriteLines = rewriteAddSizeOfARInstruction(hm, functiontable, commentStr, markerStr,	operand);	
			} else if( "LDA".equals(instruction) ) {
				rewriteLines = rewriteLDAInstruction(commentStr, markerStr,	symbolicOffset);
			} else if( "STA".equals(instruction) ) {
				rewriteLines = rewriteSTAInstruction(commentStr, markerStr,	symbolicOffset);
			} else if( "MUL".equals(instruction) ) {
				rewriteLines = rewriteMathInstruction( commentStr, markerStr, symbolicOffset, "MUL" );
			} else if( "SUB".equals(instruction) ) {
				rewriteLines = rewriteMathInstruction( commentStr, markerStr, symbolicOffset, "SUB" );
			} else if( "ADD".equals(instruction) ) {
				rewriteLines = rewriteMathInstruction( commentStr, markerStr, symbolicOffset, "ADD" );
			} else {
				// Otherwise, we just write out the line as is and hope for the best
				rewriteLines = new String[1];
				rewriteLines[0] = line;
			}
			
			// Now that we know what the replacement lines are, write them out
			for( String rewriteLine : rewriteLines ) {
				outputWriter.write( rewriteLine );
				outputWriter.write( Program.NEWLINE );
			}
		}
		
		// Let others know if we're in a new procedure
		return proc;
	}

	/**
	 * Since the functiontable hashes by the username and not the location
	 * label (i.e., marker) if we want to find one by the marker, we have to
	 * do a little more work
	 * 
	 * @param functiontable Hashed by user name
	 * @param marker A location label
	 * @return A matching proc
	 */
	private Proc findProcedureForMarker(HashMap<String, Proc> functiontable, String marker) {
		
		for( Proc proc : functiontable.values() ) {
			if( proc.getLocationLabel().equals(marker) ) {
				return proc;
			}
		}
		return null;
	}

	/**
	 * This happens when we're rewriting an math instruction, because
	 * it's actually adding from within an activation record.
	 * See slide 32. Note that since there's no '<math-operator> indirect',
	 * you will need to load it indirectly, then store it in
	 * a fixed location (like the Scratch area) that you can
	 * operate from.
	 * 
	 * @param commentStr Comments originally at the end of the line
	 * @param markerStr Markers originally at the front of the line
	 * @param symbolicOffset Its symbolic offset, guaranteed non-null
	 * @param mathInstruction The type of math operation (ADD, SUB, or MUL)
	 * @return
	 */
	private String[] rewriteMathInstruction(String commentStr, String markerStr, 
			String symbolicOffset, String mathInstruction) {
		
		String[] rewriteLines = new String[8];
		
		// Remember we have to preserve labels and comments!
		StringBuffer strbuf = new StringBuffer("");
		if( markerStr != null ) {
			strbuf.append(markerStr + " ");
		}
		// Store what's in the accumulator to scratch temporarily
		strbuf.append( "STA " + SymbolicNames.getSymbolicNameForScratch() );
		// Remember we have to preserve labels and comments!
		if( commentStr != null ) {
			strbuf.append( " ; " + commentStr );
		}
		rewriteLines[0] = strbuf.toString();
		// Scratch now has the original accumulator value: the first operand
		
		// Now we need to load the second operand into memory
		// Its value is stored at memory location FP+symbolicOffset
		rewriteLines[1] = "LDA " + SymbolicNames.getSymbolicNameForFP();
		rewriteLines[2] = "ADD " + symbolicOffset;
		//the Accumulator now has the memory address of the second operand's value
		
		rewriteLines[3] = "STA " + SymbolicNames.getSymbolicNameForFPB();
		//the FPB now has the memory address of the second operand's value
		
		// Next we load it indirectly
		rewriteLines[4] = "LDI " + SymbolicNames.getSymbolicNameForFPB();
		// Accumulator now has the value of the second operand

		// Now we need to store the second operand
		rewriteLines[5] = "STA " + SymbolicNames.getSymbolicNameForFPB();
		// The FPB now has the value of the second operand
		
		// And load the first operand
		rewriteLines[6] = "LDA " + SymbolicNames.getSymbolicNameForScratch();
		// The accumulator now has the first operand value
		
		// And perform the math operation on the second operand
		rewriteLines[7] = mathInstruction + " " + SymbolicNames.getSymbolicNameForFPB();
		// The accumulator should now have our answer
		
		// And perform the math, at long last!
		return rewriteLines;
	}

	/**
	 * This happens when we're rewriting a LDA instruction, because
	 * it's actually loading from within an activation record.
	 * See slide 32.
	 * 
	 * @param commentStr Comments originally at the end of the line
	 * @param markerStr Markers originally at the front of the line
	 * @param symbolicOffset Its symbolic offset, guaranteed non-null
	 * @return
	 */
	private String[] rewriteLDAInstruction(String commentStr, String markerStr,
			String symbolicOffset) {
		
		String[] rewriteLines = new String[4];
		
		// Remember we have to preserve labels and comments!
		StringBuffer strbuf = new StringBuffer("");
		if( markerStr != null ) {
			strbuf.append(markerStr + " ");
		}
		// First line is loading the frame pointer
		strbuf.append( "LDA " + SymbolicNames.getSymbolicNameForFP() );
		// Remember we have to preserve labels and comments!
		if( commentStr != null ) {
			strbuf.append( " ; " + commentStr );
		}
		rewriteLines[0] = strbuf.toString();
		
		// Second line is adding a symbolic offset
		rewriteLines[1] = "ADD " + symbolicOffset;
		
		// Third line is storing it in the FPB
		rewriteLines[2] = "STA " + SymbolicNames.getSymbolicNameForFPB();
		
		// Fourth line is a load indirect
		rewriteLines[3] = "LDI " + SymbolicNames.getSymbolicNameForFPB();
		
		return rewriteLines;
	}

	/**
	 * This happens when we're rewriting a OMG instruction.
	 * 
	 * A OMG is a meta-instruction that only this rewriter knows about.
	 * It says to add the size of activation record for a particular
	 * function to the accumulator.
	 * 
	 * The reason this exists is because when making a function call,
	 * we don't always know the amount of memory we need to allocate
	 * for what we're calling. This happens in recursive functions,
	 * or when calling a function that is defined further ahead 
	 * in the source code.
	 * 
	 * @param hm The symbol table
	 * @param functiontable Hashed by user name
	 * @param commentStr Comments originally at the end of the line
	 * @param markerStr Markers originally at the front of the line
	 * @param operand The function to load the activation record size for
	 * @return
	 */
	private String[] rewriteAddSizeOfARInstruction( HashMap<String,Integer> hm, Map<String,Proc> functiontable,
			String commentStr, String markerStr, String operand) {
		String[] rewriteLines = new String[1];

		// First, look up the function in the function table
		Proc proc = functiontable.get(operand);
		String symbol = "<ERROR>";	// I am feeling pessimistic right now
		
		if( proc == null ) {
			System.err.println( "Could not find a process corresponding to '" + operand + "' so it will not be allocated memory." );
			System.err.println( "This is bad and things won't work." );
		} else {
			// 1 less because the SP is at the end of the record, not one beyond it
			symbol = SymbolicNames.getSymbolicNameForInteger(hm, (proc.getSizeOfAR()-1) );
		}
		rewriteLines[0] = "ADD " + symbol;
		
		return rewriteLines;
	}

	/**
	 * This happens when we're rewriting a HAK instruction.
	 * 
	 * A HAK is a meta-instruction that only this rewriter knows about.
	 * It's like a LDA, but instead of loading from an offset relative
	 * to the address in FP, it loads it relative to an address in Scratch.
	 * 
	 * The reason this exists is because when making a function call,
	 * we have to be able to read a variable relative to one frame,
	 * and write it relative to another. If the load rewrites are always
	 * relative to the current frame pointer, then we can't do this;
	 * we have no way of going backward.
	 * 
	 * The easiest solution was to assume that the previous pointer
	 * was stored in scratch, and load relative to that.
	 * 
	 * This is a desperate hack at a late hour....
	 * 
	 * @param commentStr Comments originally at the end of the line
	 * @param markerStr Markers originally at the front of the line
	 * @param symbolicOffset Its symbolic offset, guaranteed non-null
	 * @return
	 */
	private String[] rewriteLoadAccRelativeToPrevFPInstruction(String commentStr, String markerStr,
			String symbolicOffset) {
		
		String[] rewriteLines = new String[7];
		
		// Step 1: Load the previous frame pointer memory address into the accumulator
		// It can be found at the current SP memory location, -1
		
		// Remember we have to preserve labels and comments!
		StringBuffer strbuf = new StringBuffer("");
		if( markerStr != null ) {
			strbuf.append(markerStr + " ");
		}
		// First we load the current SP
		strbuf.append( "LDA " + SymbolicNames.getSymbolicNameForSP() );
		// Remember we have to preserve labels and comments!
		if( commentStr != null ) {
			strbuf.append( " ; " + commentStr );
		}
		rewriteLines[0] = strbuf.toString();
		
		// Then we subtract 1 -- this gives us the memory address of where the prev fp's memory address is stored
		rewriteLines[1] = "SUB Constant1";  	// Cheating, but I know Constant1 will be in memory....
		rewriteLines[2] = "STA " + SymbolicNames.getSymbolicNameForScratch();	// Store this to scratch
		rewriteLines[3] = "LDI " + SymbolicNames.getSymbolicNameForScratch();	// Now the accumulator has the prev fp memory address in it
		
		// Step 2: Add a symbolic offset
		rewriteLines[4] = "ADD " + symbolicOffset;
		
		// Step 3: Store it in Scratch
		rewriteLines[5] = "STA " + SymbolicNames.getSymbolicNameForScratch();
		
		// Step 4: Load it indirectly from scratch
		rewriteLines[6] = "LDI " + SymbolicNames.getSymbolicNameForScratch();
		
		return rewriteLines;
	}

	/**
	 * This happens when we're rewriting an STA instruction, because
	 * it's actually writing to within an activation record.
	 * See slide 32.
	 * 
	 * @param commentStr Comments originally at the end of the line
	 * @param markerStr Markers originally at the front of the line
	 * @param symbolicOffset Its symbolic offset, guaranteed non-null
	 * @return
	 */
	private String[] rewriteSTAInstruction(String commentStr, String markerStr,
			String symbolicOffset) {
		
		String[] rewriteLines = new String[6];
		
		// Remember we have to preserve labels and comments!
		StringBuffer strbuf = new StringBuffer("");
		if( markerStr != null ) {
			strbuf.append(markerStr + " ");
		}
		// Save what we want to save temporarily to scratch
		strbuf.append( "STA " + SymbolicNames.getSymbolicNameForScratch() );
		// Remember we have to preserve labels and comments!
		if( commentStr != null ) {
			strbuf.append( " ; " + commentStr );
		}
		rewriteLines[0] = strbuf.toString();
		
		// Figure out the address of where we really want to save it, and put that in FPB
		rewriteLines[1] = "LDA " + SymbolicNames.getSymbolicNameForFP();
		rewriteLines[2] = "ADD " + symbolicOffset;
		rewriteLines[3] = "STA " + SymbolicNames.getSymbolicNameForFPB();
		
		// Now load what we really want to save back into memory
		rewriteLines[4] = "LDA " + SymbolicNames.getSymbolicNameForScratch();
		
		// Final line is a save indirect
		rewriteLines[5] = "STI " + SymbolicNames.getSymbolicNameForFPB();
		
		return rewriteLines;
	}

	/**
	 * We sometimes have to replace instructions for: LDA, STA,
	 * SUB, MUL, ADD
	 * 
	 * @param instruction What to consider
	 * @return True if we might have to rewrite it
	 */
	private boolean potentialRewriteNeeded(String instruction) {
		return instruction != null &&
				(instruction.equals("LDA") || instruction.equals("STA") ||
						instruction.equals("SUB") || instruction.equals("MUL") ||
						instruction.equals("ADD") );
	}
	
}
